import numpy as np
from time import time

class MTSP:
    def __init__(self,node:int, car:int, depot:int, car_depot:dict, graph:np.ndarray, num:int, mur) -> None:
        """_summary_

        Args:
            node (int): 服务节点数量
            car (int): 使用车辆数目上限
            depot (int): 车场数量
            car_depot (dict): 车辆与车场的隶属关系
            graph (np.ndarray): 图G(V,E),前(0,1,...,k-1)表示depot,(k,k+1,...,k+n-1)表示服务节点
        """
        ## 设备信息
        self.node = node
        self.car = car
        self.depot = depot

        ## 环境信息
        self.car_depot = car_depot
        self.graph = graph

        ## 关键参数
        self.population_num = num
        self.mr = mur

        ## 邻域搜索
        self.scale = None

        self.length = node + car - 1

        ## 变量
        self.individual = None
        self.fitness = None

        ## 
        self.best = 1e+9



    def _cal_path(self, c):
        car_paths = np.zeros(self.car, np.float64)

        count = 0
        for i in range(self.length):
            if(i == 0):
                if(c[i] != -1):
                    car_paths[count] += self.graph[self.car_depot[count], c[i] + self.depot]
                else:
                    count += 1

                continue

            now = c[i] + self.depot if c[i] != -1 else self.car_depot[count]
            last = c[i-1] + self.depot if c[i] != -1 else self.car_depot[count]
            if now != last:
                car_paths[count] += self.graph[last, now]

            if c[i] == -1:
                count += 1
        
        last = c[-1] + self.depot if c[-1] != -1 else self.car_depot[self.car - 1]

        if last != self.car_depot[self.car - 1]:
            car_paths[-1] += self.graph[last, self.car_depot[self.car - 1]]

        return car_paths.sum()


    def cal_fitness(self, c):
        """适应度计算函数,可扩展重写

        Args:
            c (ndarray or list): 个体

        Returns:
            double: 适应度函数
        """
        fitness = self._cal_path(c)

        return fitness


    def init_population(self):
        self.individual = np.zeros((self.population_num,self.length), np.int64)
        for k in range(self.population_num):
            count = 0
            for i in range(self.node):
                self.individual[k,i + count] = i
                if np.random.rand() < 0.2 and count < self.car - 1:
                    count += 1
                    self.individual[k, i+count] = -1

            if count < self.car - 1:
                for j in range(self.car - count - 1):
                    self.individual[k, self.node + count + j] = -1


            idxes = np.random.choice(self.length, 2, False)

            self.individual[k, idxes[0]], self.individual[k, idxes[1]] = self.individual[k, idxes[1]], self.individual[k, idxes[0]]

        self.fitness = {i:self.cal_fitness(self.individual[i]) for i in range(self.population_num)}

        self.update()
           


    def _cross(self, p1:np.ndarray, p2:np.ndarray):
        idxes = np.sort(np.random.choice(self.length, 2, replace=False))

        arc1 = p1[idxes[0] : idxes[1] + 1]

        c1 = np.copy(p1)
        c2 = np.copy(p2)

        num = len(np.where(arc1 == -1)[0])

        count = 0
        flag = 0
        for i in range(self.length):
            if p2[i] in arc1:
                if p2[i] == -1:
                    if flag < num:
                        flag += 1
                    else:
                        continue

                c1[idxes[0] + count] = p2[i]
                c2[i] = arc1[count]
                count += 1

            if count == len(arc1):
                break

        
        return c1,c2

    def cross(self, p1, p2, num):
        if(num % 2 == 1):
            print("scale must be even")
            exit(-1)

        n = int(num / 2)

        res = np.zeros((num, self.node + self.car - 1), np.int64)
        for i in range(n):
            c1,c2 = self._cross(p1, p2)
            res[2*i] = c1
            res[2*i + 1] = c2

        return res

    def mutate(self, c):
        idxes = idxes = np.sort(np.random.choice(self.length, 2, replace=False))
        c[idxes[0]],c[idxes[1]] = c[idxes[1]],c[idxes[0]]


    
    def update(self, children = None, num:int = 40):
        
        if children is not None:
            for i in range(num):
                self.fitness[self.population_num + i] = self.cal_fitness(children[i])
            self.individual = np.vstack((self.individual, children))



        items = sorted(self.fitness.items(),key=lambda x:x[1])
        idxes = np.zeros(self.population_num, np.int64)

        self.fitness = dict()
        for i in range(self.population_num):
            idxes[i] = items[i][0]
            self.fitness[i] = items[i][1]
           
        self.individual = self.individual[idxes]
        

    def step(self, num):
        ## 选择个体
        idxes = np.random.choice(self.population_num, 2, replace=False)

        ## 交叉算子(结合邻域搜索)
        children = self.cross(self.individual[idxes[0]], self.individual[idxes[1]], num)

        ## 变异
        for i in range(num):
            if np.random.rand() < self.mr:
                self.mutate(children[i])

        ## 更新种群
        self.update(children, num)

        return self.fitness[0]

    def print_log(self, step, epoch, val, t):
        if val < self.best:
            print("[{:5d}/{:5d}] --- Best UB = {:.2f} \t Total Time = {:.4f}".format(step, epoch, val, t))

            self.best  =val

    def __call__(self, epoch, num):
        print("Iteration Times is {} and scale size is {}".format(epoch, num))
        
        self.init_population()

        print("Population Initialized Done!")

        print("Start Timing --- ")
        st = time()
        self.print_log(0, epoch, self.fitness[0], 0.0)

        for i in range(epoch):
            one = self.step(num)
            et = time()
            self.print_log(i+1, epoch, one, et - st)

        print("Total time is {:.4f}".format(time() - st))
        print("Evolution Action Done!")
        print("Best UB is {}".format(self.fitness[0]))



if __name__ == "__main__":
    node = 1000
    car = 10
    depot = 3

    car_depot = {i:np.random.randint(0, depot) for i in range(car)}

    size = node+ car - 1

    graph = np.random.randint(5,20,(size,size))

    model = MTSP(node, car, depot, car_depot, graph, 300, 0.1)

    res = model(1000, 40)